import torch


def check_operator_accuracy(
    golden: torch.Tensor,
    actual: torch.Tensor,

    mare_cpu_low: float = 0.0,
    error_threshold: float = 2**-10, 
    error_threshold_accept : float = 2**-11, 
) -> bool:
    """
    检查算子输出与参考结果的精度是否满足标准
    :param golden: 参考张量（如PyTorch结果）
    :param actual: 算子输出的张量（如Triton结果）
    :param mare_cpu_low: CPU版本算子的最低MARE值
    :param error_threshold: 判定为误差超标的阈值
    :return: 是否通过精度检查
    TODO 默认error_threshold?
    """
    # 验证输入合法性
    assert actual.shape == golden.shape, \
        f"形状不匹配: 参考结果={golden.shape}, 实际输出={actual.shape}"
    assert actual.device == golden.device, \
        f"设备不匹配: 参考结果={golden.device}, 实际输出={actual.device}"
    assert actual.dtype == golden.dtype, \
        f"数据类型不匹配: 参考结果={golden.dtype}, 实际输出={actual.dtype}"

    # 基础误差计算
    diff = actual - golden
    abs_error = torch.abs(diff)
    abs_golden = torch.abs(golden)
    abs_actual = torch.abs(actual)
    # 总元素
    total_elements = golden.numel()
    # 避免除零
    eps = 1e-6  

    # 计算核心误差指标
    # 相对误差l
    rel_error = abs_error / (abs_golden + eps)  
    # 最大相对误差
    mare = torch.max(rel_error).item() if total_elements > 0 else 0.0  
    # 平均相对误差
    mere = torch.mean(rel_error).item() if total_elements > 0 else 0.0  
    # 均方根误差
    rmse = torch.sqrt(torch.mean(diff**2)).item() if total_elements > 0 else 0.0  
    # 误差均衡性(EB)
    max_abs_golden = torch.max(abs_golden).item() if total_elements > 0 else 0.0
    eb_denominator = max(max_abs_golden, 1.0)
    eb = torch.mean(diff / eb_denominator).item() if total_elements > 0 else 0.0
    # 补充：最大绝对误差、最大绝对误差
    max_abs_error = torch.max(abs_error).item() if total_elements > 0 else 0.0
    mean_abs_error = torch.mean(abs_error).item() if total_elements > 0 else 0.0

    # 相对误差率统计
    rel_error_mask = rel_error > error_threshold
    rel_error_count = torch.sum(rel_error_mask).item()
    rel_error_rate = rel_error_count / total_elements if total_elements > 0 else 0.0
    
    # 补充：绝对误差率统计
    abs_error_mask = abs_error > error_threshold
    abs_error_count = torch.sum(abs_error_mask).item()
    abs_error_rate = abs_error_count / total_elements if total_elements > 0 else 0.0

    # 定位最大相对误差坐标
    max_rel_indices = ()
    first_rel_coord = ()
    if total_elements > 0:
        # mare_tensor = torch.max(rel_error)  # 保留张量形式，避免精度损失
        # mare = mare_tensor.item()  # 仅在需要数值时转换
        # max_rel_mask = rel_error == mare_tensor  # 用张量比较
        # max_rel_indices = torch.nonzero(max_rel_mask, as_tuple = True)
        # # 转换为可读的坐标格式（取第一个最大误差点）
        # first_rel_coord = tuple(coord[0].item() for coord in max_rel_indices)
        mare_tensor = torch.max(rel_error)  # 保留张量形式，避免精度损失
        mare = mare_tensor.item()  # 仅在需要数值时转换
        max_rel_mask = rel_error == mare_tensor  # 用张量比较
        max_rel_indices = torch.nonzero(max_rel_mask, as_tuple=True)
        # 检查是否有有效坐标
        if all(coord.numel() > 0 for coord in max_rel_indices):
            first_rel_coord = tuple(coord[0].item() for coord in max_rel_indices)
        else:
            first_rel_coord = ()  # 或其他默认值
    else:
        first_rel_coord = ()    

    # 定位最大绝对误差坐标
    max_abs_indices = ()
    first_abs_coord = ()
    # if total_elements > 0: 
    #     max_abs_mask = abs_error == max_abs_error
    #     max_abs_indices = torch.nonzero(max_abs_mask, as_tuple = True)
    #     # 转换为可读的坐标格式（取第一个最大误差点）
    #     first_abs_coord = tuple(coord[0].item() for coord in max_abs_indices)
    # else: 
    #     first_abs_coord = ()
    if total_elements > 0:
    # 保留张量形式的最大值，避免精度损失
        max_abs_error_tensor = torch.max(abs_error)
        max_abs_error = max_abs_error_tensor.item()  # 仅在需要数值时转换
        max_abs_mask = abs_error == max_abs_error_tensor  # 用张量比较
        max_abs_indices = torch.nonzero(max_abs_mask, as_tuple=True)
        
        # 检查是否有有效坐标，避免访问空张量
        if all(coord.numel() > 0 for coord in max_abs_indices):
            first_abs_coord = tuple(coord[0].item() for coord in max_abs_indices)
        else:
            first_abs_coord = ()  # 或其他默认值（如(-1, -1)表示无有效坐标）
    else:
        first_abs_coord = ()


    # 打印所有统计信息
    # print("=" * 60)
    print("[Accuracy Check Result]")
    print(f"- [Tensor info] Shape = {actual.shape}, Elements num = {total_elements}, Device = {actual.device}. \n")
    
    print("[Main error metrics]")
    print(f"- [MARE(Max Reletive Error)] {mare:.8f}")
    print(f"- [MERE(Mean Reletive Error)] {mere:.8f}")
    print(f"- [RMSE(Root Mean Squared Error)] {rmse:.8f}")
    print(f"- [EB(Error Balance)] {eb:.8f}")
    # 补充：最大绝对误差、最大绝对误差
    print(f"- [Max Absolute Error] {max_abs_error:.8f}")
    print(f"- [Mean Absolute Error] {mean_abs_error:.8f}\n")
    
    print("[Error Distribution Statistics]")
    print(f"- [Error Threshold] {error_threshold:.8f}")
    print(f"- [Golden Data Range] ({torch.min(golden):.8f}, {torch.max(golden):.8f})")
    print(f"- [Actual Data Range] ({torch.min(actual):.8f}, {torch.max(actual):.8f})")
    print(f"- [Reletive Error Elements Num] {rel_error_count}")
    print(f"- [Reletive Error Rate] {rel_error_rate:.6%}")
    # 补充：绝对误差率统计
    print(f"- [Absolute Error Elements Num] {abs_error_count}")
    print(f"- [Absolute Error Rate] {abs_error_rate:.6%}\n")
    
    # print("[Max Reletive Error Info]")
    # if total_elements > 0: 
    #     print(f"- [Max Reletive Error Coord] {first_rel_coord}")
    #     print(f"- [Golden Data] {golden[first_rel_coord].item():.8f}")
    #     print(f"- [Actual Data] {actual[first_rel_coord].item():.8f}")
    #     print(f"- [Abs Error]{abs_error[first_rel_coord].item():.8f}")
    #     print(f"- [Rel Error] {rel_error[first_rel_coord].item():.8f}\n")
    print("[Max Reletive Error Info]")
    if total_elements > 0:
        # 检查坐标有效性：非空且维度与张量匹配
        is_valid = (len(first_rel_coord) == len(golden.shape) 
                    and all(0 <= coord < dim for coord, dim in zip(first_rel_coord, golden.shape)))
        
        if is_valid:
            print(f"- [Max Reletive Error Coord] {first_rel_coord}")
            print(f"- [Golden Data] {golden[first_rel_coord].item():.8f}")
            print(f"- [Actual Data] {actual[first_rel_coord].item():.8f}")
            print(f"- [Abs Error] {abs_error[first_rel_coord].item():.8f}")
            print(f"- [Rel Error] {rel_error[first_rel_coord].item():.8f}\n")
        else:
            print(f"- [Max Reletive Error Coord] Invalid or empty coordinates {first_rel_coord}")
            print(f"- [Note] No valid max relative error point found (may be due to floating-point precision)\n")

    # print("[Max Absolute Error Info]")
    # if total_elements > 0:
    #     print(f"- [Max Absolute Error Coord] {first_abs_coord}")
    #     print(f"- [Golden Data] {golden[first_abs_coord].item():.8f}")
    #     print(f"- [Actual Data] {actual[first_abs_coord].item():.8f}")
    #     print(f"- [Abs Error] {abs_error[first_abs_coord].item():.8f}")
    #     print(f"- [Rel Error] {rel_error[first_abs_coord].item():.8f}\n")
    print("[Max Absolute Error Info]")
    if total_elements > 0:
        print(f"- [Max Absolute Error Coord] {first_abs_coord}")
        # 检查坐标有效性：非空且维度与张量形状匹配
        is_valid = (len(first_abs_coord) == len(golden.shape) 
                    and all(0 <= coord < dim for coord, dim in zip(first_abs_coord, golden.shape)))
        
        if is_valid:
            print(f"- [Golden Data] {golden[first_abs_coord].item():.8f}")
            print(f"- [Actual Data] {actual[first_abs_coord].item():.8f}")
            print(f"- [Abs Error] {abs_error[first_abs_coord].item():.8f}")
            print(f"- [Rel Error] {rel_error[first_abs_coord].item():.8f}\n")
        else:
            print(f"- [Warning] Invalid coordinates {first_abs_coord}, cannot retrieve scalar value\n")
    
    
    # 验收标准检查
    denominator = max(mare_cpu_low, error_threshold_accept)
    ratio = mare / denominator
    passed = ratio < 10
    
    print(f"[Final Pass Result] MARE_npu / max(MARE_cpu_low, 2^-11 = {error_threshold_accept:.6f}) = {ratio:.6f} < 10 ? {passed}\n")

    return passed